#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <stdlib.h>
#include <stddef.h>
#include <assert.h>
#include <poll.h>
#include <errno.h>
#include <time.h>
#include <sys/time.h>
#include <sys/socket.h>
#include <sys/uio.h>
#include <netinet/in.h>

#include "rlib.h"

#define DATASIZE 500
#define ACKLEN 8
#define DATALEN 12

int verify_cksum(rel_t *r, packet_t *pkt, size_t n);
int send_ackpkt(conn_t *c, int ackno);
int check_partial(rel_t *r);
uint32_t num_ms();

/* stores a packet and relevant information about it */
typedef struct pkt_wrapper {
	packet_t *pkt;
	int valid;
	int partial;
	int timed_out;
} pkt;

struct reliable_state {
	struct sockaddr_storage *a;
	int window;
	int timeout;
	uint32_t prevtime; /* used to calculate timeout */

  rel_t *next;			/* Linked list for traversing all connections */
  rel_t **prev;
  conn_t *c;			/* This is the connection object */

	/* packets to store in window */
	pkt **sent_data;
	pkt **recv_data;
	int send_window_start;
	int recv_window_start;

	/* flags */
	int last_seg_sent; /* seqno */
	int sent_eof;
	int read_eof;

	/* temporary buffer storage */
	char buf[DATASIZE];
	int bufsize;
};

rel_t *rel_list = NULL;

/* check if a partial packet is in flight. 1 if true, 0 if false */
int check_partial(rel_t *r) {
	int i;
	for(i = 0; i < r->window; i++) {
		if(r->sent_data[i]->valid == 1 && r->sent_data[i]->partial == 1) {
			return 1;
		}
	}
	return 0;
}

int send_ackpkt(conn_t *c, int ackno) {
	/* create ack packet to send */
	struct ack_packet ackp;
	ackp.cksum = 0;
	ackp.len = htons(ACKLEN);
	ackp.ackno = htonl(ackno);
	ackp.cksum = cksum(&ackp, ACKLEN);
	return conn_sendpkt(c, (packet_t *) &ackp, ACKLEN);
}



/* Creates a new reliable protocol session, returns NULL on failure.
 * Exactly one of c and ss should be NULL.  (ss is NULL when called
 * from rlib.c, while c is NULL when this function is called from
 * rel_demux.) */
rel_t *
rel_create (conn_t *c, const struct sockaddr_storage *ss,
	    const struct config_common *cc)
{
  rel_t *r;

  r = xmalloc (sizeof (*r));
  memset (r, 0, sizeof (*r));

  if (!c) {
    c = conn_create (r, ss);
    if (!c) {
      free (r);
      return NULL;
    }
  }

	r->window = cc->window;
	r->timeout = cc->timeout;

  r->c = c;
  r->next = rel_list;
  r->prev = &rel_list;
  if (rel_list)
    rel_list->prev = &r->next;
  rel_list = r;

	/* initialize flags */
	r->a = malloc(sizeof(struct sockaddr_storage));
	if(ss != NULL) {
		memcpy(r->a, ss, sizeof(struct sockaddr_storage));
	}
	r->prevtime = 0;
	r->bufsize = 0;
	r->sent_eof = 0;
	r->read_eof = 0;
	r->send_window_start = 1;
	r->recv_window_start = 1;
	r->last_seg_sent = 0;

	/* initialize space for packets stored */
	r->sent_data = malloc(sizeof(pkt *) * r->window);
	r->recv_data = malloc(sizeof(pkt *) * r->window);

	int i;
	for(i = 0; i < r->window; i++) {
		r->sent_data[i] = malloc(sizeof(pkt));
		r->sent_data[i]->pkt = malloc(sizeof(packet_t));
		r->sent_data[i]->valid = 0;
		r->sent_data[i]->partial = 0;
		r->sent_data[i]->timed_out = 0;
		r->recv_data[i] = malloc(sizeof(pkt));
		r->recv_data[i]->pkt = malloc(sizeof(packet_t));
		r->recv_data[i]->valid = 0;
	}

  return r;
}

void
rel_destroy (rel_t *r)
{
  if (r->next)
    r->next->prev = r->prev;
  *r->prev = r->next;
  conn_destroy (r->c);

	/* free space for packets stored */
	int i;
	for(i = 0; i < r->window; i++) {
		free(r->sent_data[i]->pkt);
		free(r->sent_data[i]);
		free(r->recv_data[i]->pkt);
		free(r->recv_data[i]);
	}

	free(r->sent_data);
	free(r->recv_data);
	free(r->a);
	free(r);
}


/* This function only gets called when the process is running as a
 * server and must handle connections from multiple clients.  You have
 * to look up the rel_t structure based on the address in the
 * sockaddr_storage passed in.  If this is a new connection (sequence
 * number 1), you will need to allocate a new conn_t using rel_create
 * ().  (Pass rel_create NULL for the conn_t, so it will know to
 * allocate a new connection.)
 */
void
rel_demux (const struct config_common *cc,
	   const struct sockaddr_storage *ss,
	   packet_t *pkt, size_t len)
{
	rel_t *r = rel_list;
	while(r != NULL) {
		/* check if sockaddr matches */
		if(addreq(r->a, ss)) {
			break;
		}
		r = r->next;
	}

	if(r == NULL) {
		/* check that it is a new connection (seqno == 1). */
		if(verify_cksum(r, pkt, len) == 0 && ntohl(pkt->seqno) == 1) {
			r = rel_create(NULL, ss, cc);
			rel_recvpkt(r, pkt, len);
		} else {
			return;
		}
	} else {
		rel_recvpkt(r, pkt, len);
	}

}

/* given a packet, verify its checksum. return 0 if checksum is correct, -1 on 
 * error */
int verify_cksum(rel_t *r, packet_t *pkt, size_t n) {
	uint16_t checksum = pkt->cksum;
	uint16_t len = ntohs(pkt->len);
  
	/* check the checksum */	
	pkt->cksum = 0;
	if(len > DATALEN + DATASIZE || checksum != cksum(pkt, len)) {
		/* error with verifying checksum */
		pkt->cksum = checksum; /* revert checksum value */
		return -1;
	}
	pkt->cksum = checksum; /* revert checksum value */
	return 0;
}

void
rel_recvpkt (rel_t *r, packet_t *pkt, size_t n)
{
	uint16_t len = ntohs(pkt->len);
	if(verify_cksum(r, pkt, n) < 0 || len != n) {
		return;
	}
	uint32_t ackno = ntohl(pkt->ackno);

	if(len != ACKLEN) { /* is a data packet */
		uint32_t seqno = ntohl(pkt->seqno); /* Only valid if length > ACKLEN */
		if(seqno < r->recv_window_start) { /* duplicate packet */
			return;
		}

		int diff = seqno - r->recv_window_start + 1;
		if(diff > r->window) {
			return;
		}

		int slot = seqno - r->recv_window_start;
		memcpy(r->recv_data[slot]->pkt, pkt, sizeof(packet_t)); /* copy pkt into r->recv_data */
		r->recv_data[slot]->valid = 1;
		rel_output(r); /* send data to be outputted */
	} else {
		if(ackno <= r->send_window_start) {
			/* duplicate ack */
			return;
		}

		/* shift down entries */
		int j;
		int loop_start = ackno - r->send_window_start;
		for(j = loop_start; j < r->window; j++) {
			struct pkt_wrapper *newp = r->sent_data[j];
			struct pkt_wrapper *oldp = r->sent_data[j-loop_start];

			if(newp->valid == 0) {
				oldp->valid = 0;
				break;
			}

			/* shift down all contents of the packet_wrapper */
			memcpy(oldp->pkt, newp->pkt, sizeof(packet_t));	
			oldp->valid = newp->valid;
			oldp->partial = newp->partial;
			oldp->timed_out = 1; /* indicate timeout for packet */
		}

		while(j < r->window + loop_start) {
			/* zero out all info for slots from this point on */
			struct pkt_wrapper *oldp = r->sent_data[j-loop_start];
			oldp->valid = 0;
			oldp->partial = 0;
			oldp->timed_out = 0;
			j++;
		}
		
		r->send_window_start = ackno;
		rel_read(r);
	}
}


void
rel_read (rel_t *s)
{	
	int diff = s->last_seg_sent - s->send_window_start + 1;
	if(diff == 0 && s->sent_eof == 1 && s->read_eof == 1) {
		/* tear down connection */
		rel_destroy(s);
		return;
	}

	if(diff >= s->window || s->sent_eof) {
		/* haven't received ack yet, don't read input yet */
		return;
	}

	char temp[DATASIZE];
	int num_bytes = conn_input(s->c, temp, DATASIZE);
	
	if(num_bytes == -1) {
		/* EOF or error received */
		s->sent_eof = 1;
		num_bytes = 0;
	} else if(num_bytes == 0) {
		/* nothing to add to buffer */
		return;
	}

	char buf2[DATASIZE]; /* to store overflow bytes */
	int overflow_bytes = num_bytes - (DATASIZE - s->bufsize);
	if(overflow_bytes >= 0) {
		/* copy as much as you can into s->buf, store remainder into buf2 */
		memcpy(s->buf + s->bufsize, temp, DATASIZE - s->bufsize);
		memcpy(buf2, temp + (DATASIZE - s->bufsize), overflow_bytes);
		s->bufsize = DATASIZE;
	} else {
		/* copy all remaining bits into s->buf, update bufsize */
		memcpy(s->buf + s->bufsize, temp, num_bytes);
		s->bufsize = s->bufsize + num_bytes;
	}

	/* send the packet only if it is an EOF, if there was overflow, or if it is a 
	 * partial packet and no partial packet is currently in flight. */
	int should_send = num_bytes == 0 || overflow_bytes >= 0 ||
		(overflow_bytes < 0 && check_partial(s) == 0);
	if(!should_send) {
		return;
	}

	/* now, build the packet to send */
	s->last_seg_sent += 1;
	packet_t pkt;
	pkt.cksum = 0;
	int len = DATALEN + s->bufsize;
	pkt.len = htons(len);
	pkt.ackno = htonl(s->recv_window_start);
	pkt.seqno = htonl(s->last_seg_sent);
	memcpy(pkt.data, s->buf, s->bufsize); /* s->buf is the real data to send */
	pkt.cksum = cksum(&pkt, len);
	conn_sendpkt(s->c, &pkt, len);

	/* adjust flags, store packet just sent (for possible resending) */
	int slot = s->last_seg_sent - s->send_window_start;
	memcpy(s->sent_data[slot]->pkt, &pkt, len); /* copy pkt into r->recv_data */
	s->sent_data[slot]->valid = 1;

	/* populate s->buf with the contents of buf2 */
	if(overflow_bytes >= 0) {
		memcpy(s->buf, buf2, overflow_bytes);
		s->bufsize = overflow_bytes;
	} else if(num_bytes == 0) {
		/* EOF, do nothing special */
	} else {
		s->bufsize = 0;
		s->sent_data[slot]->partial = 1;
	}
}

void
rel_output (rel_t *r)
{
	int diff = r->last_seg_sent - r->send_window_start + 1;	

	int i;
	for(i = 0; i < r->window; i++) {
		if(r->recv_data[i]->valid == 0) {
			break;
		}

		packet_t *datapkt = r->recv_data[i]->pkt;
		uint16_t pktlen = ntohs(datapkt->len);
		size_t len = conn_bufspace(r->c);
		if(len >= pktlen - DATALEN) {
			if(pktlen == DATALEN) {
				r->read_eof = 1;
			}
			conn_output(r->c, datapkt->data, pktlen - DATALEN);
			r->recv_data[i]->valid = 0;	/* to signal that it has been output */
		} else {
			/* cannot output anymore, halt */
			break;
		}
	}

	r->recv_window_start += i; /* update seqno pointed to by window start */
	send_ackpkt(r->c, r->recv_window_start);

	if(diff == 0 && r->sent_eof == 1 && r->read_eof == 1) {
		/* tear down connection */
		rel_destroy(r);
		return;
	}

	/* shift entries down in recv list */
	int j;
	int loop_start;
	for(loop_start = 0; loop_start < r->window; loop_start++) {
		if(r->recv_data[loop_start]->valid == 1) {
			/* found point where outputting stopped */
			break;
		}
	}

	for(j = loop_start; j < r->window; j++) {
		if(r->recv_data[j - loop_start]->valid == 0) {
			break;
		}

		/* shift data down */
		memcpy(r->recv_data[j-loop_start]->pkt, r->recv_data[j]->pkt, 
				sizeof(packet_t));
		r->recv_data[j-loop_start]->valid = 1;
	}

	/* set remainder entries to be invalid */
	while(j < r->window + loop_start) {
		r->recv_data[j-loop_start]->valid = 0;
		j++;
	}
}

uint32_t num_ms() {
	uint32_t ms;
	struct timeval tv;
	gettimeofday(&tv, NULL);
	ms = tv.tv_sec * 1000 + (tv.tv_usec / 1000);
	return ms;
}

void
rel_timer ()
{
	rel_t *r = rel_list;
	uint32_t currtime = num_ms();

	while(r != NULL) {
		uint32_t difft = currtime - r->prevtime;
		
		if(difft < r->timeout) {
			r = r->next;
			continue; /* not time yet */
		}

		r->prevtime = currtime;
		int i;
		int sent_single = 0; /* only send first non-requested packet */
		for(i = 0; i < r->window; i++) {
			pkt *p = r->sent_data[i];
							
			if(p->valid == 0) {
				continue;
			}

			int resend = 0; /* boolean to determine if we should resend packet */
			if(p->timed_out == 1) {
				resend = 1;
				p->timed_out = 0; /* un-time-out the packet */
			} else if(sent_single == 0) {
				resend = 1;
				sent_single = 1; /* don't send any more non-timed-out packets */
			}

			if(resend) {
				conn_sendpkt(r->c, p->pkt, ntohs(p->pkt->len));
			}
		}

		/* find first packet that has not been received (in order */
		for(i = 0; i < r->window; i++) {
			if(r->recv_data[i]->valid == 0) {
				break;
			}
		}
		
		/* create ack packet to send */
		if(r->recv_window_start + i > 1) {
			send_ackpkt(r->c, r->recv_window_start + i);
		}

		r = r->next;
	}
}
